'''
Created on 

This module defines the functions and classes that support
the VM Assisted Puppet Tier 3 Installation. 
https://twiki.cern.ch/twiki/bin/view/Atlas/VMPuppetAssistedTier3Install


Classes

AT3Node - A Tier 3 Node description
AT3Cfg - The Tier 3 Cluster configuration description
Role - Describes a role played by a node


Functions

AT3_GenerateAllKS - generates the kickstart scripts
AT3_GenerateEtcHosts - generates the etchosts file
AT3_GenerateNodesPP - generates the puppet congfiguration script nodes.pp
AT3_GeneratePXELinuxCFG - generates the PXE linux configuration file
AT3_GenerateVMCreationScripts - generates the virtual machine scripts
getIPList - returns a list of ip addresses


@author: yyao@lbl.gov, vchendrix@lbl.gov
'''
import sys
import os

class AT3Node:
    def __init__(self,role,hostname, publicip, publicnetmask, 
            privateip, privatenetmask, gateway, dns,kstemplate):
        self.role=role
        self.hostname=hostname
        self.publicip=publicip
        self.privateip=privateip
        self.gateway=gateway
        self.dns=dns
        self.publicnetmask=publicnetmask
        self.privatenetmask=privatenetmask
        self.kstemplate=kstemplate

    def __repr__(self):
        return "(role=%s,hostname=%s,publicip=%s,publicnetmask=%s,privateip=%s,privatenetmask=%s,gateway=%s,dns=%s,kstemplate=%s)"%(self.role,str(self.hostname),str(self.publicip),str(self.publicnetmask),str(self.privateip),str(self.privatenetmask),str(self.gateway),str(self.dns),str(self.kstemplate))

class Role:
    def __init__(self,name,kstemplate, privateip,haspublicip,privatedns,privategateway,unique,privatenetmask,virtual=False):
    
        self.name=name
        self.kstemplate=kstemplate

        self.haspublicip=haspublicip

        self.privateip=privateip
        self.privatedns=privatedns
        self.privategateway=privategateway
        self.privatenetmask=privatenetmask

        self.unique=unique
        self.virtual=virtual
        
    def __repr__(self):
        return "name=%s, kstemplate=%s, haspublicip=%s, privateip=%s, privatedns=%s, privategateway=%s, privatenetmask=%s, unique=%s"%(self.name,self.kstemplate,self.haspublicip,self.privateip,self.privatedns,self.privategateway,self.privatenetmask,self.unique)
        
class AT3Cfg:
    def __init__(self,name='my',privateSubnetPrefix="192.168.100",privateNetmask="255.255.255.0",
            nfsIpSuffix="1",headIpSuffix="2",ldapIpSuffix="3",proxyIpSuffix="4",
            intIpSuffix="30",workerIpSuffix="60"):
         
        path = os.getcwd()
        self.cfgDir = "./%s" % name
        self.templatesDir = "%s/templates" % self.cfgDir

        self.name=name
        self.nodes={}
        privateDNSAndGateway="%s.%s" % (privateSubnetPrefix,nfsIpSuffix)
        # Map role name to tuple(kstemplate,InternalIP/StartIP,PublicIPIsNone,
        # PrivateDNS,PrivateGateWay,unique,privateNetMask)
        self.rolelist={
            # All Roles have to point DNS to the NFS node, except NFS node itself, 
            # which will point to the outside DNS
            "NFS": Role("NFS", "ks_nfs.cfg.template", "%s.%s" %
                (privateSubnetPrefix,nfsIpSuffix) ,True,None,None,True,privateNetmask),
            "HEAD": Role("HEAD", "ks_head.cfg.template","%s.%s"
                % (privateSubnetPrefix,headIpSuffix) ,True,privateDNSAndGateway, 
                privateDNSAndGateway,True,privateNetmask),
            "LDAP": Role("LDAP", "ks_ldap.cfg.template","%s.%s" %
                (privateSubnetPrefix,ldapIpSuffix) ,False,privateDNSAndGateway,
                privateDNSAndGateway,True,privateNetmask,True),
            "PROXY": Role("PROXY", "ks_proxy.cfg.template","%s.%s" %
                (privateSubnetPrefix,proxyIpSuffix) ,False,privateDNSAndGateway,
                privateDNSAndGateway,True,privateNetmask,True),
            "INTERACTIVE": Role("INTERACTIVE", "ks_interactive.cfg.template","%s.%s" %
                (privateSubnetPrefix,intIpSuffix) ,True,privateDNSAndGateway,
                privateDNSAndGateway,False,"255.255.255.0"),
            "WORKER": Role("WORKER", "ks_worker.cfg.template","%s.%s"
                % (privateSubnetPrefix,workerIpSuffix) ,False,privateDNSAndGateway,
                privateDNSAndGateway,False,"255.255.255.0",True),
            }
        self.cfglist={}

        

    def setServiceNode_2Nic(self,role,hostname,publicip,netmask,gateway=None,dns=None,overwrite=False):
        if not role in self.rolelist.keys():
            raise RuntimeError("Role %s not in the predefined role list %s"%(str(role),str(self.rolelist.keys())))

        if not self.rolelist[role].unique:
            raise RuntimeError("Use the addXXX functions to add %s roles"%role)

        if self.nodes.has_key(role) and len(self.nodes[role])>0 and not overwrite:
            raise RuntimeError("Role %s already defined, use overwrite=True to overwrite"%(str(role)))

        if not self.rolelist[role].haspublicip:
            raise RuntimeError("Role %s has no public IP, call setServiceNode_1Nic instead"%str(role))
            
        if (publicip is None or gateway is None or dns is None) and role=="NFS":
            raise RuntimeError("Role %s has public IP, must provide publicip,gateway and dns"%str(role))

        m_publicip=publicip

        #only use outside dns for NFS server, everyone else point dns to NFS server. 
        if not self.rolelist[role].privatedns is None:
            m_dns=self.rolelist[role].privatedns
        else:
            m_dns=dns
        if not self.rolelist[role].privategateway is None:
            m_gateway=self.rolelist[role].privategateway
        else:
            m_gateway=gateway

        m_privateip=self.rolelist[role].privateip
        m_privatenetmask=self.rolelist[role].privatenetmask
        
        self.nodes[role]=[AT3Node(role=role,hostname=hostname,publicip=m_publicip,publicnetmask=netmask,privateip=m_privateip,privatenetmask=m_privatenetmask,gateway=m_gateway, dns=m_dns, kstemplate=self.rolelist[role].kstemplate)]

    def setServiceNode_1Nic(self,role,hostname,overwrite=False):
        if not role in self.rolelist.keys():
            raise RuntimeError("Role %s not in the predefined role list %s"%(str(role),str(self.rolelist.keys())))

        if not self.rolelist[role].unique:
            raise RuntimeError("Use the addXXX functions to add %s roles"%role)

        if self.nodes.has_key(role) and len(self.nodes[role])>0 and not overwrite:
            raise RuntimeError("Role %s already defined, use overwrite=True to overwrite"%(str(role)))

        if self.rolelist[role].haspublicip:
            raise RuntimeError("Role %s has public IP, call setServiceNode_2Nic instead"%str(role))
            
        m_publicip=None
        m_gateway=self.rolelist[role].privategateway
        m_dns=self.rolelist[role].privatedns
        m_privateip=self.rolelist[role].privateip
        m_privatenetmask=self.rolelist[role].privatenetmask

        self.nodes[role]=[AT3Node(role=role,hostname=hostname,publicip=m_publicip,publicnetmask=None,privateip=m_privateip,privatenetmask=m_privatenetmask,gateway=m_gateway, dns=m_dns, kstemplate=self.rolelist[role].kstemplate)]

    def addWorkerNode(self, hostname):
        role="WORKER"

        if not self.nodes.has_key(role):
            self.nodes[role]=[]

        m_publicip=None
        m_gateway=self.rolelist[role].privategateway
        m_dns=self.rolelist[role].privatedns
        m_privateip=self.rolelist[role].privateip
        m_privatenetmask=self.rolelist[role].privatenetmask
        
        self.nodes[role].append(AT3Node(role=role,hostname=hostname,publicip=m_publicip,publicnetmask=None,privateip=m_privateip,privatenetmask=m_privatenetmask,gateway=m_gateway, dns=m_dns, kstemplate=self.rolelist[role].kstemplate))

        #inc the ip address
        self.rolelist[role].privateip=self.nextip(m_privateip,1)

    def addInteractiveNode(self,hostname,publicip,netmask,gateway=None,dns=None):
        role="INTERACTIVE"

        if not self.nodes.has_key(role):
            self.nodes[role]=[]

        m_publicip=publicip
        m_gateway=self.rolelist[role].privategateway
        m_dns=self.rolelist[role].privatedns
        m_privateip=self.rolelist[role].privateip
        m_privatenetmask=self.rolelist[role].privatenetmask
        
        self.nodes[role].append(AT3Node(role=role,hostname=hostname,publicip=m_publicip,publicnetmask=netmask,privateip=m_privateip,privatenetmask=m_privatenetmask,gateway=m_gateway, dns=m_dns, kstemplate=self.rolelist[role].kstemplate))

        #inc the ip address
        self.rolelist[role].privateip=self.nextip(m_privateip,1)
        
    def set(self,name,value):
        self.cfglist[name]=value

    def get(self,name):
        if not self.cfglist.has_key(name):
            raise RuntimeError("configuration with name=%s does not exist"%name)
        return self.cfglist[name]
        
    def nextip(self,ipstarts,n):
        spip=ipstarts.split(".")
        spip[3]=str(int(spip[3])+n)
        if int(spip[3])>=255:
            raise RuntimeError("IP address out of range %s + %s"%(ipstarts,str(n)))
        return ".".join(spip)

        
def AT3_GenerateAllKS(cfg):

    print "Generating kickstart file for: ",
    
    for irole in cfg.rolelist:
        nodes=cfg.nodes[irole]
        for inode in nodes:
            #build the list of variables
            rlist={}
            for ikey in cfg.cfglist.keys():
                rlist[ikey]=cfg.cfglist[ikey]

            rlist["PUBLICIP"]=inode.publicip
            rlist["PUBLICNETMASK"]=inode.publicnetmask            

            rlist["PRIVATEIP"]=inode.privateip
            rlist["PRIVATENETMASK"]=inode.privatenetmask            

            rlist["HOSTNAME"]=inode.hostname

            rlist["GATEWAY"]=inode.gateway
            rlist["DNS"]=inode.dns

            #            print rlist

            print '%s '%inode.hostname.split(".")[0],
            
            #open template
            inf=open("%s/%s" % (cfg.templatesDir,inode.kstemplate),"r")
            tmpstr=inf.read()
            inf.close()

            #write output
            ouf=open("%s/kickstart-%s.cfg"%(cfg.cfgDir, inode.hostname.split(".")[0]),"w")
            ouf.write(tmpstr%rlist)
            ouf.close()

    print
            
    
def AT3_GenerateEtcHosts(cfg):
    print "Generating /etc/hosts file for NFS node"
    
    stretchosts=""

    #localhost stuff
    stretchosts=stretchosts+"127.0.0.1 localhost.localdomain localhost\n"


    for iRole in cfg.nodes.keys():
        for iNode in cfg.nodes[iRole]:
            if iRole=="NFS":
                stretchosts=stretchosts+"%s puppet %s\n"%(iNode.privateip,iNode.hostname.split(".")[0])
            else:
                stretchosts=stretchosts+"%s %s\n"%(iNode.privateip,iNode.hostname.split(".")[0]) 

    ouf=open("%s/etchosts" % cfg.cfgDir,"w")
    ouf.write(stretchosts)
    ouf.close()

    
    
def AT3_GenerateVMCreationScripts(cfg):


    print "Generating vm creation script for: ",
    
    ksserver=cfg.nodes["NFS"][0].privateip
    
    for irole in cfg.rolelist:
        if not cfg.rolelist[irole].virtual:
            continue
        nodes=cfg.nodes[irole]
        for inode in nodes:
            rlist={}
            for ikey in cfg.cfglist.keys():
                rlist[ikey]=cfg.cfglist[ikey]

            name=inode.hostname.split(".")[0]
            rlist["VMNAME"]="vm_%s"%name
            rlist["DISKNAME"]="vmdisk_%s"%name

            rlist["KS_SERVER"]="http://%s:8080"%ksserver
            rlist["KS_NAME"]="kickstart-%s.cfg"%name
            
            rlist["PRIVATEIP"]=inode.privateip
            rlist["PRIVATENETMASK"]=inode.privatenetmask            

            rlist["HOSTNAME"]=inode.hostname
            rlist["GATEWAY"]=inode.gateway
            rlist["DNS"]=inode.dns

            #            print rlist

            print '%s '%name,

            #open template
            inf=open("%s/createvm.sh.template" % cfg.templatesDir,"r")
            tmpstr=inf.read()
            inf.close()


            #write output
            ouf=open("%s/crvm-%s.sh"%( cfg.cfgDir,name), "w")
            ouf.write(tmpstr%rlist)
            ouf.close()

    print
    
    
def AT3_GeneratePXELinuxCFG(cfg):

    print "Generating pxe boot menu item for:",
    
    nodes=cfg.nodes["WORKER"]

    mastertemplate="""
default menu.c32
prompt 0
timeout 300
ONTIMEOUT local

MENU TITLE PXE Menu

LABEL local
 MENU LABEL Boot local hard drive
 LOCALBOOT 0

%(ITEMS)s
"""
    itemtemplate="""
LABEL %(SHORTHOSTNAME)s
 MENU LABEL %(SHORTHOSTNAME)s
 KERNEL images/vmlinuz
 APPEND initrd=images/initrd.img ramdisk_size=100000 url --url %(URL)s ksdevice=eth0 gateway=%(GATEWAY)s dns=%(DNS)s netmask=%(PRIVATENETMASK)s ip=%(PRIVATEIP)s ks=%(KS_SERVER)s/kickstart-%(SHORTHOSTNAME)s.cfg
"""
    itemstr=""

    ksserver=cfg.nodes["NFS"][0].privateip
    for inode in nodes:
        #build the list of variables
        rlist={}
        for ikey in cfg.cfglist.keys():
            rlist[ikey]=cfg.cfglist[ikey]
            
        rlist["PRIVATEIP"]=inode.privateip
        rlist["PRIVATENETMASK"]=inode.privatenetmask            
        rlist["HOSTNAME"]=inode.hostname
        rlist["SHORTHOSTNAME"]=inode.hostname.split(".")[0]
        rlist["GATEWAY"]=inode.gateway
        rlist["DNS"]=inode.dns
        rlist["KS_SERVER"]="http://%s:8080"%ksserver
        
        #            print rlist
        print '%s '%inode.hostname.split(".")[0],

        itemstr=itemstr+itemtemplate%rlist
        
    #write output
    ouf=open("%s/pxelinux.cfg.default" %cfg.cfgDir,"w")
    ouf.write(mastertemplate%{"ITEMS":itemstr})
    ouf.close()
        
    print

def getIPList(cfg,roles):
    #Get a list of IP address of all the given roles, in the private network only
    iplist=[]
    for irole in roles:
        for inode in cfg.nodes[irole]:
            iplist.append(inode.privateip)
    return "\""+"\",\"".join(iplist)+"\""

def AT3_GenerateNodesPP(cfg):

    print "Generating nodes.pp entry for:",

    outputstr=""

    for irole in cfg.rolelist:

        nodes=cfg.nodes[irole]
        shortldapname=cfg.nodes["LDAP"][0].hostname.split(".")[0]
        for inode in nodes:
            rlist={}

            shorthostname=inode.hostname.split(".")[0]

            print shorthostname,

            inf=open("%s/nodes.template.%s.pp"%(cfg.templatesDir, irole),"r")
            tmpinput=inf.read()
            inf.close()

            #IP and Hostnames
            rlist["SHORTHOSTNAME"]=shorthostname
            rlist["FQDN"]=inode.hostname
            rlist["DNS"]=inode.dns
            rlist["DOMAINNAME"]=cfg.get("DOMAINNAME")
            rlist["PUBLICNETMASK"]=inode.publicnetmask
            rlist["PUBLICIP"]=inode.publicip
            rlist["GATEWAY"]=inode.gateway
            rlist["PRIVATENETMASK"]=inode.privatenetmask
            rlist["PRIVATEIP"]=inode.privateip
            rlist["PUBLICNETMASK"]=inode.publicnetmask
            rlist["PUBLICIP"]=inode.publicip

            #LDAP stuff
            rlist["LDAPHOSTNAME"]=shortldapname
            rlist["LDAPSUFFIX"]=cfg.get("LDAPSUFFIX")
            rlist["LDAPROOTPW"]=cfg.get("LDAPROOTPW")

            #Ganglia Stuff
            rlist["CLUSTERNAME"]=cfg.get("CLUSTERNAME")
            rlist["GANGLIACLIENTLIST"]=getIPList(cfg,["NFS","HEAD","LDAP","PROXY","INTERACTIVE","WORKER"])
            rlist["GANGLIASRVADDR"]=cfg.nodes["NFS"][0].privateip

            #NFS Stuff
            rlist["NFSCLIENTLIST"]=cfg.get("NFSCLIENTLIST")
            rlist["NFSHOST"]=cfg.nodes["NFS"][0].privateip

            #Condor Stuff
            rlist["CONDORFILESYSTEMDOMAIN"]=cfg.get("CONDORFILESYSTEMDOMAIN")
            rlist["CONDORPASSWORD"]=cfg.get("CONDORPASSWORD")
            rlist["CONDORHEADADDR"]=cfg.nodes["HEAD"][0].privateip

            #For xrootd
            rlist["HEADHOSTNAME"]=cfg.nodes["HEAD"][0].hostname

            #Proxy Stuff
            rlist["PROXYNAME"]=cfg.get("PROXYNAME")
            rlist["PROXYCACHESIZE"]=cfg.get("PROXYCACHESIZE")
            rlist["PROXYEMAIL"]=cfg.get("PROXYEMAIL")
            rlist["PROXYCACHEMEM"]=cfg.get("PROXYCACHEMEM")

            rlist["SQUIDPROXY"]="http://%s:3128"%cfg.nodes["PROXY"][0].hostname
            
            #Other
            rlist["URL"]=cfg.get("URL")
            rlist["PRIVATESUBNETPREFIX"]=cfg.get("PRIVATESUBNETPREFIX")

            tmpoutput=tmpinput%rlist
            outputstr=outputstr+tmpoutput

    print

    ouf=open("%s/nodes.pp" % cfg.cfgDir,"w")
    ouf.write(outputstr)
    ouf.close()
    

